{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import Caltech101\n",
    "from torchvision import models\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision import datasets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载预训练的 ConvNeXt-Tiny 模型\n",
    "# model = models.convnext_tiny(pretrained=True)\n",
    "model = models.convnext_tiny(weights=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # 处理验证集数据\n",
    "# import os\n",
    "# import shutil\n",
    "# # 当前目录\n",
    "# val_dir = '/root/dr/CI/data_tiny_imagenet/tiny-imagenet-200/val'\n",
    "# val_images_dir = os.path.join(val_dir, 'images')\n",
    "# val_annotations_file = os.path.join(val_dir, 'val_annotations.txt')\n",
    "# # 创建用于分类的文件夹结构\n",
    "# with open(val_annotations_file, 'r') as f:\n",
    "#     for line in f.readlines():\n",
    "#         parts = line.split('\\t')\n",
    "#         img_file = parts[0]\n",
    "#         class_id = parts[1]\n",
    "#         # 创建类别文件夹\n",
    "#         class_dir = os.path.join(val_dir, class_id)\n",
    "#         if not os.path.exists(class_dir):\n",
    "#             os.makedirs(class_dir)   \n",
    "#         # 移动图像到类别文件夹\n",
    "#         shutil.move(os.path.join(val_images_dir, img_file), os.path.join(class_dir, img_file))\n",
    "# # 删除空的 images 文件夹\n",
    "# os.rmdir(val_images_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据预处理和增强\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomResizedCrop((224, 224)),  # 随机裁剪并调整大小到 224x224\n",
    "    transforms.RandomHorizontalFlip(),  # 随机水平翻转\n",
    "    transforms.ToTensor(),  # 转换为张量\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化\n",
    "])\n",
    "\n",
    "transform_val = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),  # 调整大小到 224x224\n",
    "    transforms.ToTensor(),  # 转换为张量\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化\n",
    "])\n",
    "\n",
    "data_dir = '/root/dr/CI/data_tiny_imagenet/tiny-imagenet-200/'\n",
    "# 加载训练集\n",
    "train_dataset = datasets.ImageFolder(root=data_dir+'train/', transform=transform_train)\n",
    "train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)\n",
    "\n",
    "# 加载验证集\n",
    "val_dataset = datasets.ImageFolder(root=data_dir+'val/', transform=transform_val)\n",
    "val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 替换最后的分类层以适应 tiny_imagenet 的 200 个类别\n",
    " # ConvNeXt-Tiny 的最后一层输出是 768 维\n",
    "model.classifier[2] = torch.nn.Linear(in_features=768, out_features=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on device: cuda\n"
     ]
    }
   ],
   "source": [
    "# 将模型移到 GPU（如果可用）\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = model.to(device)\n",
    "print(f'Training on device: {device}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "# 记录训练过程中的损失和精度\n",
    "train_losses = []\n",
    "train_accuracies = []\n",
    "val_accuracies = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义计算准确率的函数\n",
    "def calculate_accuracy(loader, model):\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    model.eval()  # 切换到评估模式\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in loader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            outputs = model(inputs)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100, Train Loss: 5.1589, Train Acc: 0.0177, Val Acc: 0.0312\n",
      "Epoch 2/100, Train Loss: 4.8585, Train Acc: 0.0369, Val Acc: 0.0576\n",
      "Epoch 3/100, Train Loss: 4.6388, Train Acc: 0.0608, Val Acc: 0.0804\n",
      "Epoch 4/100, Train Loss: 4.4718, Train Acc: 0.0832, Val Acc: 0.1165\n",
      "Epoch 5/100, Train Loss: 4.2613, Train Acc: 0.1098, Val Acc: 0.1452\n",
      "Epoch 6/100, Train Loss: 4.0553, Train Acc: 0.1385, Val Acc: 0.1818\n",
      "Epoch 7/100, Train Loss: 3.8940, Train Acc: 0.1629, Val Acc: 0.2073\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 100 # 训练 100 个 epoch\n",
    "# 训练模型（简单示例）\n",
    "for epoch in range(num_epochs):\n",
    "    if epoch == 20:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = 0.005\n",
    "        print(f\"Learning rate adjusted to {optimizer.param_groups[0]['lr']} at epoch {epoch}\")\n",
    "    if epoch == 60:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = 0.0001\n",
    "        print(f\"Learning rate adjusted to {optimizer.param_groups[0]['lr']} at epoch {epoch}\")\n",
    "    if epoch == 80:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = 0.00002\n",
    "        print(f\"Learning rate adjusted to {optimizer.param_groups[0]['lr']} at epoch {epoch}\")\n",
    "    model.train()  # 切换到训练模式\n",
    "    running_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for inputs, labels in train_loader:\n",
    "        inputs, labels = inputs.to(device), labels.to(device)\n",
    "\n",
    "        # 前向传播\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # 反向传播和优化\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        # 计算训练准确率\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "    train_loss = running_loss / len(train_loader)\n",
    "    train_acc = correct / total\n",
    "    val_acc = calculate_accuracy(val_loader, model)\n",
    "\n",
    "    train_losses.append(train_loss)\n",
    "    train_accuracies.append(train_acc)\n",
    "    val_accuracies.append(val_acc)\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 绘制训练损失、训练精度和验证精度的曲线图\n",
    "epochs_range = range(1, num_epochs + 1)\n",
    "plt.figure(figsize=(10, 5))\n",
    "\n",
    "# 绘制训练损失\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(epochs_range, train_losses, label='Train Loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Train Loss')\n",
    "plt.legend()\n",
    "\n",
    "# 绘制训练精度和验证精度\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(epochs_range, train_accuracies, label='Train Accuracy')\n",
    "plt.plot(epochs_range, val_accuracies, label='Val Accuracy')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Train and Val Accuracy')\n",
    "plt.legend()\n",
    "\n",
    "# 保存图像到工作目录\n",
    "plt.tight_layout()\n",
    "plt.savefig('training_results.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存模型的 state_dict 到磁盘\n",
    "torch.save(model.state_dict(), 'convnext_tinyimagenet_100epochs.pth')\n",
    "\n",
    "print(\"模型已保存为 convnext_100epochs.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
